// iMSI-Compiler.cpp : Defines the entry point for the console application.
//

#include "stdafx.h"
#include "../Common/Compression/Compress.h"

#ifdef _UNICODE
typedef std::wstring String;
#else
typedef std::string String;
#endif

_TCHAR* Source = 0;
_TCHAR* Target = 0;
_TCHAR* Engine = 0;

_TCHAR PathToEngine[MAX_PATH] = {0};

struct FileInfo
{
    String Name;
    size_t Offset;
    size_t Length;
    FileInfo ()
        : Offset (0), Length (0)
    { }
};

class OSException : public std::exception
{
public:
	typedef exception _Mybase;

	explicit OSException (const String& aMessage)
		: _Mybase(aMessage.c_str())
	{  }

	explicit OSException (const char* aMessage)
		: _Mybase(aMessage)
	{  }
};

class ScopedHandle
{
private:
    HANDLE Item;
    ScopedHandle (const ScopedHandle&);
    void operator= (const ScopedHandle&);
public:
    ScopedHandle (HANDLE aItem)
        : Item (aItem)
    {  }

    ~ScopedHandle ()
    {
        if (Item != INVALID_HANDLE_VALUE) {
            CloseHandle (Item);
        }
        Item = INVALID_HANDLE_VALUE;
    }

    HANDLE operator() () const
    {
        return Item;
    }
};

void RaiseOsError (const char* aMessage)
{
    DWORD Error = GetLastError ();
    if (Error != ERROR_SUCCESS) {
        char Buffer[2048] = {0};
        strcpy (Buffer, aMessage);
        strcat (Buffer, "\n");
        FormatMessageA (FORMAT_MESSAGE_FROM_SYSTEM, 0, Error, 1, Buffer + strlen (aMessage), 2048, 0);
        throw OSException (Buffer);
    }
}

void PrintUsage()
{
    _tprintf(TEXT("Usage:\n"));
    _tprintf(TEXT("\timc <Source directory> <Setup filename>\n"));
}

bool DirectoryExists (_TCHAR* aPath)
{
    DWORD Attributes = GetFileAttributes(aPath);
    return (Attributes != INVALID_FILE_ATTRIBUTES && (Attributes & FILE_ATTRIBUTE_DIRECTORY));
}

void SetupPathToEngine()
{
    _TCHAR Self[MAX_PATH] = {0};

    if (GetModuleFileName(NULL, Self, MAX_PATH)) {
        TCHAR Drive[_MAX_DRIVE] = {0};
        TCHAR Directory[_MAX_DIR] = {0};
        TCHAR FileName[_MAX_FNAME] = {0};
        TCHAR Extension[_MAX_EXT] = {0};
        _tsplitpath(Self, Drive, Directory, FileName, Extension);
        _tcscat (Directory, TEXT("\\Engine"));

        _tmakepath (PathToEngine, Drive, Directory, 0, 0);

        if (DirectoryExists(PathToEngine)) {
            Engine = PathToEngine;
        }
    }
}

void AnalyseCommandLine(int argc, _TCHAR* argv[])
{
    if (argc < 3 || !DirectoryExists (argv[1])) {
        PrintUsage();
        return;
    }

    Source = argv[1];
    Target = argv[2];

    SetupPathToEngine();

    if (Engine == 0) {
        _tprintf(TEXT("Unable to find iMSI engine directory - exit\n"));
        return;
    }
}

void EnsureOutputDirectory()
{
    TCHAR Drive[_MAX_DRIVE] = {0};
    TCHAR Directory[_MAX_DIR] = {0};
    TCHAR FileName[_MAX_FNAME] = {0};
    TCHAR Extension[_MAX_EXT] = {0};
    _tsplitpath(Target, Drive, Directory, FileName, Extension);

    _TCHAR FinalDirectory[MAX_PATH] = {0};
    _tmakepath (FinalDirectory, Drive, 0, 0, 0);

    _TCHAR* Single = _tcstok(Directory, TEXT("\\"));

    while (Single != 0) {
        _tcscat (FinalDirectory, TEXT("\\"));
        _tcscat (FinalDirectory, Single);

        if (CreateDirectory(FinalDirectory, NULL) || GetLastError () == ERROR_ALREADY_EXISTS) {
            Single = _tcstok(NULL, TEXT("\\"));
        }
        else {
            RaiseOsError ("EnsureOutputDirectory - creating directory");
        }
    }
}

void PlaceEngine ()
{
    _TCHAR SourceName[MAX_PATH] = {0};
    _tcscpy (SourceName, Engine);
    _tcscat (SourceName, TEXT("iMSI-Engine.exe"));

    if (!CopyFile (SourceName, Target, FALSE)) {
        String ErrorInfo = TEXT("PlaceEngine - CopyFile\n");
        ErrorInfo += SourceName;
        ErrorInfo += TEXT(" -> ");
        ErrorInfo += Target;
        RaiseOsError (ErrorInfo.c_str ());
    }
}

void ExtractFileName (const _TCHAR* aFullPath, _TCHAR* aOutput)
{
    TCHAR Extension[_MAX_EXT] = {0};
    _tsplitpath(aFullPath, 0, 0, aOutput, Extension);
    _tcscat(aOutput, Extension);
}

void CompressFile(const _TCHAR* aFullName)
{
    String Target = aFullName;
    Target += TEXT(".compressed");
    if (EncodeFile(aFullName, Target.c_str()) != 0) {
        throw std::runtime_error("Error compressing file");
    }
}

size_t AppendFile (HANDLE aOutput, _TCHAR* aFullName)
{
    size_t Result = 0;
    CompressFile(aFullName);
    String Source = aFullName;
    Source += ".compressed";
    ScopedHandle Input = CreateFile(Source.c_str(), GENERIC_READ, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);
    if (Input () == INVALID_HANDLE_VALUE) {
        RaiseOsError ("AppendFile - opening input");
    }

    unsigned char Buffer[2048];
    DWORD Read;
    DWORD Written;
    do {
        if (!ReadFile(Input (), Buffer, 2048, &Read, 0)) {
            RaiseOsError ("AppendFile - reading input");
        }
        if (!WriteFile(aOutput, Buffer, Read, &Written, 0)) {
            RaiseOsError ("AppendFile - writing output");
        }
        Result += Written;
    } while (Read > 0);

    return Result;
}

void CompileFiles(HANDLE aOutput, std::list<FileInfo>& aFiles, _TCHAR* aSource, bool aSkipEngine)
{
    _TCHAR SearchPattern[MAX_PATH] = {0};
    _tcscpy(SearchPattern, aSource);
    _tcscat(SearchPattern, TEXT("\\*.*"));

    WIN32_FIND_DATA Data;

    HANDLE Search = FindFirstFile(SearchPattern, &Data);

    if (Search != INVALID_HANDLE_VALUE) {
        do {
            if ((Data.dwFileAttributes & FILE_ATTRIBUTE_DIRECTORY) == 0) {
                _TCHAR FileName[MAX_PATH] = {0};
                ExtractFileName (Data.cFileName, FileName);

                if (_tcsicmp(FileName, TEXT("iMSI-Engine.exe")) != 0 || !aSkipEngine) {
                    FileInfo Info;
                    Info.Name = FileName;
                    Info.Offset = GetFileSize(aOutput, 0);

                    _TCHAR FullPath[MAX_PATH] = {0};
                    _tcscpy(FullPath, aSource);
                    _tcscat(FullPath, TEXT("\\"));
                    _tcscat(FullPath, Data.cFileName);

                    Info.Length = AppendFile (aOutput, FullPath);

                    if (Info.Length > 0) {
                        aFiles.push_back(Info);
                    }
                }
            }
        } while (FindNextFile(Search, &Data));

        FindClose(Search);
    }
}

void CompileEngineFiles(HANDLE aOutput, std::list<FileInfo>& aFiles)
{
    CompileFiles(aOutput, aFiles, Engine, true);
}

void CompileUserFiles(HANDLE aOutput, std::list<FileInfo>& aFiles)
{
    CompileFiles(aOutput, aFiles, Source, false);
}

void CompileSetup()
{
    PlaceEngine();
    ScopedHandle Output = CreateFile(Target, GENERIC_READ | GENERIC_WRITE, 0, 0, OPEN_EXISTING, FILE_ATTRIBUTE_NORMAL, 0);

    if (Output () == INVALID_HANDLE_VALUE) {
        RaiseOsError ("CompileSetup - Opening output");
    }

    std::list<FileInfo> Files;
    SetFilePointer(Output (), 0, 0, FILE_END);
    CompileEngineFiles(Output (), Files);
    CompileUserFiles(Output (), Files);

    DWORD First = GetFileSize(Output (), 0);

    std::list<FileInfo>::iterator Start = Files.begin();
    std::list<FileInfo>::iterator Eol = Files.end();

    DWORD Written;
    while (Start != Eol) {
		_TCHAR Buffer[2048] = {0};
		_stprintf(Buffer, TEXT("\"%hs\"%.10d %.10d"), Start->Name.c_str(), Start->Offset, Start->Length);
        if (!WriteFile(Output (), Buffer, _tcslen(Buffer), &Written, 0)) {
            RaiseOsError ("CompileSetup - writing output");
        }
        ++Start;
    }

	_TCHAR Buffer[2048] = {0};
	_stprintf(Buffer, TEXT(" %.10d %.10d %hs"), Files.size(), First, TEXT("OLJA"));
    if (!WriteFile(Output (), Buffer, _tcslen (Buffer), &Written, 0)) {
        RaiseOsError ("CompileSetup - writing output");
    }
}

int _tmain(int argc, _TCHAR* argv[])
{
	try {
        AnalyseCommandLine(argc, argv);

        if (Source == 0 || Target == 0 || Engine == 0) {
            return 1;
        }

        EnsureOutputDirectory();
	    CompileSetup();

        return 0;
	}
	catch (std::exception& aError) {
        printf (aError.what ());
		return 1;
	}
}
